package org.orz.cloud.common.utils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;


public class ThreadLocalUtils {
    private static final ThreadLocal<Map<Object, Object>> threadBucket = new ThreadLocalUtils.InheritableThreadLocalMap<>();
    private static final Logger log = LoggerFactory.getLogger(ThreadLocalUtils.class);

    public ThreadLocalUtils() {
    }

    public static <T> T getValue(Object key, Class<T> clazz) {
        Map<Object, Object> bucket = threadBucket.get();
        if (null == bucket) {
            if (log.isWarnEnabled()) {
                log.warn("ThreadLocal bucket: " + key + " not exists.");
            }
            return null;
        } else {
            Object valueSrc = bucket.get(key);
            if (clazz.isInstance(valueSrc)) {
                return (T) valueSrc;
            } else {
                if (log.isWarnEnabled()) {
                    if (null == valueSrc) {
                        log.warn("Cannot find object from ThreadLocal bucket: " + key + "; no object found");
                    } else {
                        log.warn("Cannot find object from ThreadLocal bucket: " + key + "; except type is " + clazz.getSimpleName() + " but got " + valueSrc.getClass().getSimpleName());
                    }
                }
                return null;
            }
        }
    }

    public static <T> void setValue(Object key, T value) {
        Map<Object, Object> bucket = threadBucket.get();
        if (null == bucket) {
            Map<Object, Object> newBucket = new HashMap<>();
            threadBucket.set(newBucket);
            bucket = newBucket;
        }

        bucket.put(key, value);
    }

    public static void clear(Object key) {
        Map<Object, Object> bucket = threadBucket.get();
        if (null != bucket) {
            bucket.remove(key);
        }
    }

    public static void clear() {
        threadBucket.remove();
    }

    private static final class InheritableThreadLocalMap<T extends Map<Object, Object>> extends InheritableThreadLocal<Map<Object, Object>> {
        private InheritableThreadLocalMap() {
        }

        protected Map<Object, Object> childValue(Map<Object, Object> parentValue) {
            return parentValue != null ? (Map<Object, Object>)((HashMap)parentValue).clone() : null;
        }
    }
}
